from runner import Runner
from common.arguments import get_args
from common.utils import make_encirclement_env
import matplotlib.pyplot as plt
import numpy as np
import torch
import random

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def set_random_seed(seed):
    """设置所有随机种子以确保结果可复现"""
    # Python内置random模块
    random.seed(seed)
    # NumPy随机种子
    np.random.seed(seed)
    # PyTorch随机种子
    torch.manual_seed(seed)
    # 如果使用CUDA
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        # 确保CUDA操作的确定性
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    print(f"随机种子已设置为: {seed}")

def plot_encirclement_environment(env):
    """绘制合围环境 - 支持多目标"""
    grid_size = env.grid_size
    
    fig, ax = plt.subplots(figsize=(12, 10))
    
    # 绘制所有目标
    target_colors = ['red', 'darkred', 'crimson', 'maroon']
    for i, target in enumerate(env.multi_target.targets):
        target_pos = target['position']
        color = target_colors[i % len(target_colors)]
        ax.plot(target_pos[0], target_pos[1], 's', color=color, markersize=15, 
               label=f'目标{i+1}')
        
        # 绘制目标的合围范围
        circle_min = plt.Circle(target_pos, env.min_encirclement_distance, 
                               color=color, fill=False, linestyle='--', alpha=0.3)
        circle_max = plt.Circle(target_pos, env.max_encirclement_distance, 
                               color=color, fill=False, linestyle='--', alpha=0.3)
        ax.add_artist(circle_min)
        ax.add_artist(circle_max)
    
    # 绘制无人机
    colors = ['blue', 'green', 'orange', 'purple', 'brown', 'pink']
    for i, uav in enumerate(env.uavs):
        color = colors[i % len(colors)]
        ax.plot(uav.position[0], uav.position[1], 'o', color=color, 
               markersize=10, label=f'无人机{uav.id}')
        
        # 绘制无人机到最近目标的连线
        min_distance = float('inf')
        closest_target_pos = None
        for target in env.multi_target.targets:
            target_pos = target['position']
            distance = np.sqrt((uav.position[0] - target_pos[0])**2 + 
                             (uav.position[1] - target_pos[1])**2)
            if distance < min_distance:
                min_distance = distance
                closest_target_pos = target_pos
        
        if closest_target_pos:
            ax.plot([uav.position[0], closest_target_pos[0]], 
                   [uav.position[1], closest_target_pos[1]], 
                   color=color, alpha=0.3, linestyle=':')
    
    # 设置坐标轴范围
    ax.set_xlim(0, grid_size)
    ax.set_ylim(0, grid_size)
    
    # 添加图例和标题
    ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    ax.set_title('五无人机追逐双动态目标任务环境')
    ax.set_xlabel('X坐标')
    ax.set_ylabel('Y坐标')
    ax.grid(True, alpha=0.3)
    
    # 显示合围状态
    if env.check_encirclement():
        ax.text(10, grid_size-20, '成功合围目标!', fontsize=16, 
               bbox=dict(boxstyle="round", facecolor='lightgreen'))
    else:
        progress = env.calculate_encirclement_progress()
        ax.text(10, grid_size-20, f'追逐进行中... (进度: {progress:.1%})', fontsize=16,
               bbox=dict(boxstyle="round", facecolor='lightyellow'))
    
    # 显示目标数量和无人机数量信息
    ax.text(10, grid_size-50, f'无人机数量: {len(env.uavs)}', fontsize=12,
           bbox=dict(boxstyle="round", facecolor='lightblue'))
    ax.text(10, grid_size-80, f'目标数量: {len(env.multi_target.targets)}', fontsize=12,
           bbox=dict(boxstyle="round", facecolor='lightcoral'))
    
    plt.show()

if __name__ == '__main__':
    # 获取参数
    args = get_args()
    
    # 设置随机种子
    set_random_seed(args.seed)
    env, args = make_encirclement_env(args)
    # 可视化初始环境
    # plot_encirclement_environment(env)
    
    runner = Runner(args, env)
    ber = 0
    runner.bit_flip_probability = ber

    if args.evaluate:
        returns = runner.evaluate()
        print('Average returns is', returns)
    else:

        runner.encirclement_train()  # 使用新的训练方法




